
import sys
import os
import math

sys.path.append("/archive10/bowtie")

import sampleParams

def sampleMean(aa):
    return sum(aa)/(1.0*len(aa))

def sampleVariance(aa):
    mu = sampleMean(aa)
    return sum([(x-mu)*(x-mu) for x in aa])/(1.0*len(aa)-1.0)

def pearsonCorrelation(aa, bb):
    n = len(aa)
    mu_1 = sampleMean(aa)
    mu_2 = sampleMean(bb)
    var_1 = sampleVariance(aa)
    var_2 = sampleVariance(bb)
    return sum([(aa[i]-mu_1)*(bb[i]-mu_2) for i in xrange(n)]) / ((n-1.0)*math.sqrt(1.0*var_1*var_2))

def averageLinkageDistance(pts, pts2=None):
    sm = 0.0
    if pts2==None:
        points = pts
    else:
        points = pts[:]
        points.extend(pts2)
    for pt1 in points:
        for pt2 in points:
            sm += pt1.distance(pt2)
    if len(points)<=1:
        return 0.0
    return sm / (len(points)*(len(points)-1.0))

class SamplePoint:
    def __init__(self, name, samples):
        self.name = name
        self.samples = samples[:]

    def distance(self, other):
        return 1.0 - pearsonCorrelation(self.samples, other.samples)

class SampleClustering:
    def __init__(self, points):
        self.points = points[:]
        self.nextClusterId = 1
        self.clusters = []
        for pt in self.points:
            self.clusters.append((averageLinkageDistance([pt]), [pt], pt.name, self.nextClusterId))
            self.nextClusterId += 1
        self.clusters.sort()
        self.cachedDistances = {}
        
    def mergeClusters(self, i, j):
        (d1, pts1, struct1, id1) = self.clusters[i]
        (d2, pts2, struct2, id2) = self.clusters[j]
        if self.cachedDistances.has_key((id1, id2)):
            newDist = self.cachedDistances[(id1, id2)]
        else:
            newDist = averageLinkageDistance(pts1, pts2)
            self.cachedDistances[(id1, id2)] = newDist
        newCluster = (newDist, pts1 + pts2, [newDist, struct1, struct2], self.nextClusterId)
        self.nextClusterId += 1
        tmp = self.clusters
        self.clusters = [newCluster]
        for k in xrange(len(tmp)):
            if k!=i and k!=j:
                self.clusters.append(tmp[k])

    def getNextCluster(self):
        results = []
        for i in xrange(len(self.clusters)):
            for j in xrange(i+1,len(self.clusters)):
                (d1, pts1, struct1, id1) = self.clusters[i]
                (d2, pts2, struct2, id2) = self.clusters[j]
                if self.cachedDistances.has_key((id1, id2)):
                    newDist = self.cachedDistances[(id1, id2)]
                else:
                    newDist = averageLinkageDistance(pts1, pts2)
                    self.cachedDistances[(id1, id2)] = newDist
                results.append((newDist, i, j))
        results.sort()
        return results

def welch_t_test(aa, bb):
    n_1 = len(aa)
    n_2 = len(bb)
    mu_1 = sampleMean(aa)
    mu_2 = sampleMean(bb)
    var_1 = sampleVariance(aa)
    var_2 = sampleVariance(bb)
    if mu_1 == mu_2:
        return (0.0, 0.0)
    tval = (mu_1 - mu_2) / math.sqrt(var_1 / (1.0*n_1) + var_2 / (1.0*n_2))
    tdof = math.pow(var_1 / (1.0*n_1) + var_2 / (1.0*n_2), 2.0) / (var_1*var_1 / (1.0*n_1*n_1*(n_1-1)) + var_2*var_2 / (1.0*n_2*n_2*(n_2-1)))
    return (tval, tdof)

def getBenignAndTumorValues(gene, exons, sampleList):
    benignValues = []
    benignSamples = []
    tumorValues = []
    tumorSamples = []
    vals = [exonExpressionMap[(gene, exonNum)] for exonNum in exons]
    for ix in xrange(len(sampleList)):
        (type, sample) = sampleList[ix]
        if type=="Benign":
            benignSamples.append(sample)
            benignValues.append(sampleMean([valEntry[ix] for valEntry in vals]))
        else:
            tumorSamples.append(sample)
            tumorValues.append(sampleMean([valEntry[ix] for valEntry in vals]))
    return (benignSamples, benignValues, tumorSamples, tumorValues)

if __name__=="__main__":
    dir = "/archive10/bowtie/Expression/"
    paths = {}
    for filename in os.listdir(dir):
        if not filename.endswith(".normalized-gene-summary.txt"): continue
        sample = filename.replace(".normalized-gene-summary.txt", "")
        paths[sample] = dir+filename
    print >> sys.stderr, "# Samples:", paths.keys()
    params = sampleParams.getCurrentSampleParams()
    exonExpressionMap = {}
    sampleList = []
    for sample in paths.keys():
        path = paths[sample]
        if not params.has_key(sample):
            print >> sys.stderr, "# Skipping sample", sample, ": no params."
            continue
        if params[sample].getProperty("TissueType") != "Prostate":
            print >> sys.stderr, "# Skipping sample", sample, ": tissue != Prostate."
            continue
        if params[sample].getProperty("Treatment") != "FALSE":
            print >> sys.stderr, "# Skipping sample", sample, ": treatment != FALSE."
            continue
        if params[sample].getProperty("CellLineTissue") != "Tissue":
            print >> sys.stderr, "# Skipping sample", sample, ": not tissue."
            continue
        progression = params[sample].getProperty("Progression")
        if progression.startswith("Tumor"):
            type = "Tumor"
        elif progression.startswith("Benign"):
            type = "Benign"
        else:
            print >> sys.stderr, "# Skipping sample", sample, ": progression (%s) not Tumor or Benign." % progression
            continue
        print >> sys.stderr, "# Reading", sample, "expression from file", path
        f = open(path, "r")
        nexon = 0
        ngene = 0
        usedSample = False
        for line in f:
            cols = line.strip().split("\t")
            if cols[1]=="Exon" or cols[1]=="GeneExpression":
                usedSample = True
                if cols[1]=="Exon":
                    nexon += 1
                    key = (cols[0], cols[2])
                    val = float(cols[11])
                else:
                    ngene += 1
                    key = (cols[0], "Gene")
                    val = float(cols[12])
                if not exonExpressionMap.has_key(key):
                    exonExpressionMap[key] = []
                exonExpressionMap[key].append(val)
        if usedSample:
            sampleList.append((type, sample))
        f.close()
        print >> sys.stderr, "# Read %ld exons and %ld genes from %s" % (nexon, ngene, path)
        n1 = len(filter(lambda(x):x[0]=="Benign", sampleList))
        n2 = len(sampleList) - n1
        print >> sys.stderr, "# Data size =", len(exonExpressionMap), ", %d benign and %d cancer samples." % (n1, n2)
        if n1+n2>=1000:
            print >> sys.stderr, "# CONTINUING to analysis with >= 10 samples"
            break
    genes = {}
    printedHeader = False
    k = exonExpressionMap.keys()
    k.sort()
    for (gene, exonNum) in k:
        (benignSamples, benignValues, tumorSamples, tumorValues) = getBenignAndTumorValues(gene, [exonNum], sampleList)
        if sum(benignValues + tumorValues) == 0.0:
            continue
        (tval, tdof) = welch_t_test(benignValues, tumorValues)
        if not printedHeader:
            printedHeader = True
            print "Gene\tSubsetType\tNumber\tT-Statistic\tDegreesOfFreedom\tInvFoldChange\t%s\t%s" % ("\t".join(benignSamples), "\t".join(tumorSamples))
        if sampleMean(tumorValues)==0.0 and sampleMean(benignValues)==0.0:
            foldChange = 1.0
        else:
            foldChange = min(sampleMean(benignValues), sampleMean(tumorValues))/max(sampleMean(benignValues), sampleMean(tumorValues))
        type = "Exon"
        if exonNum=="Gene":
            type = "Gene"
        print "%s\t%s\t%s\t%f\t%f\t%f\t%s\t%s" % (gene, type, exonNum, tval, tdof, foldChange, "\t".join([str(x) for x in benignValues]), "\t".join([str(x) for x in tumorValues]))
        if not genes.has_key(gene):
            genes[gene] = []
        if exonNum=="Gene":
            continue
        genes[gene].append(("%s-Exon-%s" % (gene, exonNum), benignValues + tumorValues))

    for gene in genes.keys():
        print "# Clustering %d exons in %s" % (len(genes[gene]), gene)
        pts = []
        for (name, samples) in genes[gene]:
            pt = SamplePoint(name, samples)
            pts.append(pt)
        clust = SampleClustering(pts)
        # print clust.clusters
        while len(clust.clusters)>1:
            nxt = clust.getNextCluster()
            (d, i, j) = nxt[0]
            if d>0.5:
                break
            clust.mergeClusters(i, j)
            # print clust.clusters
        for i in xrange(len(clust.clusters)):
            print "Cluster %d of %d for %s : " % (i+1, len(clust.clusters), gene), clust.clusters[i][2]
            exonNums = [x.name.split("-")[-1] for x in clust.clusters[i][1]]
            (benignSamples, benignValues, tumorSamples, tumorValues) = getBenignAndTumorValues(gene, exonNums, sampleList)
            if sum(benignValues + tumorValues) == 0.0:
                continue
            (tval, tdof) = welch_t_test(benignValues, tumorValues)
            if sampleMean(tumorValues)==0.0 and sampleMean(benignValues)==0.0:
                foldChange = 1.0
            else:
                foldChange = min(sampleMean(benignValues), sampleMean(tumorValues))/max(sampleMean(benignValues), sampleMean(tumorValues))
            print "%s\tCluster\t%d\t%f\t%f\t%f\t%s\t%s" % (gene, i+1, tval, tdof, foldChange, "\t".join([str(x) for x in benignValues]), "\t".join([str(x) for x in tumorValues]))
